import sys
sys.path.append("..")

import gdal
import osr
import gdalnumeric
import numpy as np
import netCDF4 as nc
import xarray as xr
import utils.raster as raster
import datetime
import os
from multiprocessing import Process

def getnetCDFInfobyName(in_filename):
    nc_ds = gdal.Open(in_filename)
    if nc_ds is None:
        print('Open %s failed' % (in_filename))
        sys.exit()

    nc_info = {}

    # if the netcdf file has subdatasets
    if nc_ds.GetSubDatasets():
        for subDataset in nc_ds.GetSubDatasets():
            nc_ds_sub = gdal.Open(subDataset[0])
            nc_info.setdefault(subDataset[0].split(':')[-1],
                               []).append(getSubDatasetInfo(nc_ds_sub))
            nc_ds_sub = None

    else:
        nc_info.setdefault(getVariableFromMeta(nc_ds.GetMetadata_Dict()),
                           []).append(getSubDatasetInfo(nc_ds))

    nc_ds = None
    return nc_info


def getVariableFromMeta(metadata):
    for key, value in metadata.items():
        if 'scale_factor' in key:
            return key.split('#')[0]


def getSubDatasetInfo(subDataset):

    projection = osr.SpatialReference()
    result = projection.ImportFromWkt(subDataset.GetProjectionRef())
    # if the dataset has not contian the SpatialReference
    # specific the SpatialReference as WGS84

    if result != 0:
        projection.SetWellKnownGeogCS('WGS84')

    return [
        subDataset.GetRasterBand(1).GetNoDataValue(),
        subDataset.GetGeoTransform(), projection
    ]


def covertNetCDFToGeoTiff(in_filename, out_root, istep):
    step = istep

    if step % 2 != 0 and step > 1:
        raise ValueError("Step must be an even number or 1!!!")
        sys.exit(1)

    nc_info = getnetCDFInfobyName(in_filename)
    nc_ds = nc.Dataset(in_filename)
    # nc_ds=gdal.Open(in_filename)

    if not 'time' in nc_ds.variables:
        raise ValueError('Can not found the time variable')
        sys.exit(1)

    data_time = nc.num2date(nc_ds.variables['time'][:],
                            nc_ds.variables['time'].units)

    for variable, variable_info in nc_info.items():

        if variable == 'tcc':
            step = 4
        else:
            step = istep

        for time in range(0, (len(nc_ds.variables[variable]) - step + 1),
                          step):
            # out_file = ''
            out_path = os.path.join(out_root,
                                    datetime.datetime.strftime(data_time[time],
                                                               '%Y'), variable)

            isExists(out_path)
            out_file = '%s/%s_%s.tif' % (
                    out_path, variable, datetime.datetime.strftime(
                        data_time[time], '%Y%j'))

            if variable == 'tcc':            
                out_data = np.mean(nc_ds[variable][time:time + step], axis=0)
            elif variable == "tp":
                out_data = np.sum(nc_ds[variable][time:time + step], axis=0) * 1000.0  #meter -> millimeter
            else:
                out_data = np.mean(nc_ds[variable][time:time + step], axis=0)


            raster.save(out_file, out_data, nc_info[variable][0][0],nc_info[variable][0][1], nc_info[variable][0][2])

            if variable == 't2m':
                out_path = os.path.join(out_root, datetime.datetime.strftime(data_time[time],'%Y'), "m2m")
                isExists(out_path)
                max_temp = np.max(nc_ds[variable],axis = 0)
                out_temp = '%s/%s_%s.tif' % (out_path,'m2m',datetime.datetime.strftime(data_time[time], '%Y%j'))
                raster.save(out_temp, max_temp, nc_info[variable][0][0],nc_info[variable][0][1], nc_info[variable][0][2])

                # out_data = np.max(
                #     nc_ds[variable][int(time):int(time + (step / 2))],
                #     axis=0) + np.max(nc_ds[variable][int(time + (step / 2)):
                #                                      int(time + step)],
                #                      axis=0)

                # debug
                # np.savetxt('numpy_data.txt',nc_ds[variable][time])
                # sys.exit(1)

def isExists(out_file):
    if not os.path.exists(out_file):
        try:
            os.makedirs(out_file)
        except Exception as e:
            print(e)


def run(in_root, out_root, step, isMeasured=False):
    for root, dirs, files in os.walk(in_root):
        if len(files) > 0:
            for file in files:
                if file.endswith(".nc"):
                    covertNetCDFToGeoTiff(os.path.join(root, file), out_root, step)


if __name__ == "__main__":
    # covertNetCDFToGeoTiff("/home/hugo/Document/ssrd.nc",
    #                       '/home/hugo/Document', 8)
    # run('/run/media/hugo/Hugo Zhang/ERA_P/Add/',
    #     '/run/media/hugo/Hugo Zhang/ERA_P/add_re/',2)
    run("/mnt/fire/Original/tp","/mnt/fire/ERA",2)
    # p1=Process(target=run,args=('/run/media/hugo/Hugo Data/ERA/tcc_final/','/run/media/hugo/Hugo Zhang/ERA/',2))
    # # p2=Process(target=run,args=('/run/media/hugo/Hugo Data/ERAData/ss_8/','/run/media/hugo/Hugo Zhang/ERA/',8))
    # # p3=Process(target=run,args=('/run/media/hugo/Hugo Data/ERAData/tcc/','/run/media/hugo/Hugo Zhang/ERA/',4))
    # p1.start()
    # p2.start()
    # p3.start()
    # run('//run/media/hugo/Hugo Data/ERAData/tcc/',
    #     '/run/media/hugo/Hugo Data/Result/', 4)
